import os
import json
from itertools import combinations
from thermo import Chemical
from thermo.group_contribution.group_contribution_base import smarts_fragment_priority
from thermo.group_contribution.ppr78 import EPPR78_GROUPS_LIST, PPR78_kij
from chemicals.identifiers import dippr_compounds, sorted_CAS_key

# Load DIPPR compounds
dippr_cmpds = dippr_compounds()
# Or just use the compounds used in making the mehod
dippr_cmpds = ['n-BUTANE', '2,2-DIMETHYLBUTANE', '2,3-DIMETHYLBUTANE', 'NEOPENTANE', 'ETHANE', 'n-HEXANE', 'ISOBUTANE', 'ISOPENTANE', 'METHANE', '2-METHYLPENTANE', '3-METHYLPENTANE', 'n-PENTANE', 'PROPANE', '2,2-DIMETHYLHEXANE', '2,3-DIMETHYLHEXANE', '2,4-DIMETHYLHEXANE', '2,5-DIMETHYLHEXANE', '3,3-DIMETHYLHEXANE', '3,4-DIMETHYLHEXANE', '2,2-DIMETHYLPENTANE', '2,3-DIMETHYLPENTANE', '2,4-DIMETHYLPENTANE', '3,3-DIMETHYLPENTANE', '3-ETHYLHEXANE', '3-ETHYLPENTANE', 'n-HEPTANE', '2-METHYL-3-ETHYLPENTANE', '3-METHYL-3-ETHYLPENTANE', '2-METHYLHEPTANE', '3-METHYLHEPTANE', '4-METHYLHEPTANE', '2-METHYLHEXANE', '3-METHYLHEXANE', 'n-OCTANE', '2,2,3,3-TETRAMETHYLBUTANE', '2,2,3-TRIMETHYLBUTANE', '2,2,3-TRIMETHYLPENTANE', '2,2,4-TRIMETHYLPENTANE', '2,3,3-TRIMETHYLPENTANE', '2,3,4-TRIMETHYLPENTANE', '3,3-DIETHYLPENTANE', '2,2-DIMETHYL-3-ETHYLPENTANE', '2,4-DIMETHYL-3-ETHYLPENTANE', '2,2-DIMETHYLHEPTANE', '2,6-DIMETHYLHEPTANE', '3-ETHYLHEPTANE', '2-METHYLOCTANE', '3-METHYLOCTANE', '4-METHYLOCTANE', 'n-NONANE', '2,2,3,3-TETRAMETHYLPENTANE', '2,2,3,4-TETRAMETHYLPENTANE', '2,2,4,4-TETRAMETHYLPENTANE', '2,3,3,4-TETRAMETHYLPENTANE', '2,2,5-TRIMETHYLHEXANE', '2,4,4-TRIMETHYLHEXANE', 'n-DECANE', '2-METHYLNONANE', '3-METHYLNONANE', '4-METHYLNONANE', '5-METHYLNONANE', '3,3,5-TRIMETHYLHEPTANE', '2,2,3,3-TETRAMETHYLHEXANE', '2,2-DIMETHYLOCTANE', '2,3-DIMETHYLOCTANE', '2,4-DIMETHYLOCTANE', '2,5-DIMETHYLOCTANE', '2,6-DIMETHYLOCTANE', '2,7-DIMETHYLOCTANE', 'n-DODECANE', 'n-TETRADECANE', 'n-TRIDECANE', 'n-UNDECANE', 'n-DOCOSANE', 'n-EICOSANE', 'n-HENEICOSANE', 'n-HEPTADECANE', 'n-HEXADECANE', 'n-NONADECANE', 'n-OCTADECANE', 'n-PENTADECANE', 'n-TETRACOSANE', 'n-TRICOSANE', 'n-PENTACOSANE', 'n-HEXACOSANE', 'n-HEPTACOSANE', 'n-OCTACOSANE', 'n-NONACOSANE', 'n-TRIACONTANE', 'SQUALANE', 'n-DOTRIACONTANE', 'n-HEXATRIACONTANE', '2,2,5,5-TETRAMETHYLHEXANE', '2,2,4,4,6,8,8-HEPTAMETHYLNONANE', '3-METHYLUNDECANE', 'CYCLOPROPANE', 'CYCLOBUTANE', 'n-BUTYLCYCLOPENTANE', 'CYCLOPENTANE', '1,1-DIMETHYLCYCLOPENTANE', 'cis-1,2-DIMETHYLCYCLOPENTANE', 'trans-1,2-DIMETHYLCYCLOPENTANE', 'cis-1,3-DIMETHYLCYCLOPENTANE', 'trans-1,3-DIMETHYLCYCLOPENTANE', 'ETHYLCYCLOPENTANE', 'ISOPROPYLCYCLOPENTANE', 'METHYLCYCLOPENTANE', '1-METHYL-1-ETHYLCYCLOPENTANE', 'n-PROPYLCYCLOPENTANE', 'ISOPROPYLCYCLOHEXANE', 'CYCLOHEXANE', 'n-PROPYLCYCLOHEXANE', 'n-BUTYLCYCLOHEXANE', 'sec-BUTYLCYCLOHEXANE', 'tert-BUTYLCYCLOHEXANE', '1,1-DIMETHYLCYCLOHEXANE', 'cis-1,2-DIMETHYLCYCLOHEXANE', 'trans-1,2-DIMETHYLCYCLOHEXANE', 'cis-1,3-DIMETHYLCYCLOHEXANE', 'trans-1,3-DIMETHYLCYCLOHEXANE', 'cis-1,4-DIMETHYLCYCLOHEXANE', 'trans-1,4-DIMETHYLCYCLOHEXANE', 'ETHYLCYCLOHEXANE', 'METHYLCYCLOHEXANE', '1-trans-3,5-TRIMETHYLCYCLOHEXANE', 'n-DECYLCYCLOHEXANE', 'CYCLOHEPTANE', 'CYCLOOCTANE', 'cis-DECAHYDRONAPHTHALENE', 'trans-DECAHYDRONAPHTHALENE', 'DIAMANTANE', 'alpha-PINENE', 'beta-PINENE', 'ADAMANTANE', 'BICYCLOHEXYL', '1,2,3,4-TETRAMETHYLCYCLOHEXANE', '1,1-DIETHYLCYCLOHEXANE', 'trans-1,4-DIETHYLCYCLOHEXANE', '1,3-DIMETHYLADAMANTANE', '1-BUTENE', 'cis-2-BUTENE', 'trans-2-BUTENE', '2,3-DIMETHYL-1-BUTENE', '2,3-DIMETHYL-2-BUTENE', '3,3-DIMETHYL-1-BUTENE', '2-ETHYL-1-BUTENE', 'ETHYLENE', '1-HEXENE', 'cis-2-HEXENE', 'trans-2-HEXENE', 'cis-3-HEXENE', 'trans-3-HEXENE', 'ISOBUTENE', '2-METHYL-1-BUTENE', '3-METHYL-1-BUTENE', '2-METHYL-2-BUTENE', '2-METHYL-1-PENTENE', '3-METHYL-1-PENTENE', '4-METHYL-1-PENTENE', '4-METHYL-cis-2-PENTENE', '4-METHYL-trans-2-PENTENE', '2-METHYL-2-PENTENE', '3-METHYL-cis-2-PENTENE', '3-METHYL-trans-2-PENTENE', '1-PENTENE', 'cis-2-PENTENE', 'trans-2-PENTENE', 'PROPYLENE', '1-HEPTENE', 'cis-3-HEPTENE', 'cis-2-HEPTENE', 'trans-2-HEPTENE', '2-ETHYL-1-PENTENE', '4-METHYL-1-HEXENE', '3-METHYL-1-HEXENE', '2-METHYL-1-HEXENE', '3-ETHYL-1-PENTENE', 'trans-3-HEPTENE', '2,3,3-TRIMETHYL-1-BUTENE', '2-ETHYL-1-HEXENE', 'cis-2-OCTENE', 'cis-3-OCTENE', 'cis-4-OCTENE', '1-OCTENE', 'trans-2-OCTENE', 'trans-3-OCTENE', 'trans-4-OCTENE', '2,4,4-TRIMETHYL-1-PENTENE', '2,4,4-TRIMETHYL-2-PENTENE', '2,3-DIMETHYL-1-HEXENE', '2-METHYL-1-HEPTENE', '1-DECENE', '1-DODECENE', '1-HEXADECENE', '1-HEPTADECENE', '1-NONENE', '1-TRIDECENE', '1-TETRADECENE', '1-PENTADECENE', '1-NONADECENE', '1-EICOSENE', '1-OCTADECENE', '1-UNDECENE', 'CYCLOOCTENE', '1,2-BUTADIENE', '1,3-BUTADIENE', 'trans,trans-2,4-HEXADIENE', '1,5-HEXADIENE', '1,2-HEXADIENE', '1,4-HEXADIENE', 'ISOPRENE', '3-METHYL-1,2-BUTADIENE', 'cis-1,3-PENTADIENE', '1,2-PENTADIENE', 'trans-1,3-PENTADIENE', '1,4-PENTADIENE', '2,3-PENTADIENE', 'PROPADIENE', 'cis,trans-2,4-HEXADIENE', '2,3-DIMETHYL-1,3-BUTADIENE', '2,5-DIMETHYL-1,5-HEXADIENE', '2,5-DIMETHYL-2,4-HEXADIENE', '1,5-CYCLOOCTADIENE', '1,3-CYCLOHEXADIENE', '1,4-CYCLOHEXADIENE', 'CYCLOHEPTENE', 'CYCLOHEXENE', 'CYCLOPENTADIENE', 'DICYCLOPENTADIENE', 'METHYLCYCLOPENTADIENE', 'CYCLOPENTENE', '1-METHYLCYCLOPENTENE', '3-METHYLCYCLOPENTENE', '4-METHYLCYCLOPENTENE', 'VINYLCYCLOHEXENE', '2-NORBORNENE', 'CAMPHENE', 'alpha-PHELLANDRENE', 'beta-PHELLANDRENE', 'alpha-TERPINENE', 'gamma-TERPINENE', 'd-LIMONENE', 'TERPINOLENE', '1-PHENYLINDENE', 'VINYLNORBORNENE', '1-TRIACONTENE', '3-ETHYL-1-HEXENE', '4-METHYL-1-HEPTENE', '1-TETRACONTENE', 'cis-2-NONENE', 'trans-2-NONENE', '6-METHYL-1-HEPTENE', 'trans-2-EICOSENE', 'trans-2-PENTADECENE', '3-METHYL-1,4-PENTADIENE', '1,5,9-CYCLODODECATRIENE', 'trans-1,3-HEXADIENE', 'trans-2-METHYL-1,3-PENTADIENE', '1,9-DECADIENE', '1,3,5,7-CYCLOOCTATETRAENE', 'METHYLNORBORNENE', 'ETHYLNORBORNENE', '5-ETHYLIDENE-2-NORBORNENE', '5-METHYL-1-HEXENE', '2-METHYL-1-OCTENE', '1-METHYL-4-VINYLCYCLOHEXENE', '7-METHYL-1-OCTENE', '2-METHYL-1-NONENE', '8-METHYL-1-NONENE', 'cis-2-DECENE', 'trans-2-DECENE', 'cis-2-DODECENE', 'trans-2-DODECENE', 'BENZENE', 'CUMENE', 'ETHYLBENZENE', 'alpha-METHYLSTYRENE', 'o-METHYLSTYRENE', 'cis-1-PROPENYLBENZENE', 'n-PROPYLBENZENE', 'm-METHYLSTYRENE', 'STYRENE', 'TOLUENE', '1,2,3-TRIMETHYLBENZENE', '1,2,4-TRIMETHYLBENZENE', 'MESITYLENE', 'm-XYLENE', 'o-XYLENE', 'p-XYLENE', '1,2,3,4-TETRAMETHYLBENZENE', 'p-ISOPROPENYLSTYRENE', 'n-BUTYLBENZENE', 'CYCLOHEXYLBENZENE', 'm-DIETHYLBENZENE', 'o-DIETHYLBENZENE', 'DIPHENYLMETHANE', '1-METHYL-2-n-PROPYLBENZENE', '1-METHYL-3-n-PROPYLBENZENE', '1-METHYL-4-n-PROPYLBENZENE', 'sec-BUTYLBENZENE', '5-ETHYL-m-XYLENE', '3-ETHYL-o-XYLENE', '4-ETHYL-o-XYLENE', '2-ETHYL-m-XYLENE', '2-ETHYL-p-XYLENE', '4-ETHYL-m-XYLENE', 'tert-BUTYLBENZENE', 'n-DODECYLBENZENE', '1,1-DIPHENYLETHANE', 'm-CYMENE', 'n-DECYLBENZENE', 'o-CYMENE', 'n-TRIDECYLBENZENE', 'p-CYMENE', 'TRIPHENYLETHYLENE', 'p-DIETHYLBENZENE', 'n-HEXYLBENZENE', 'n-HEPTYLBENZENE', 'n-OCTYLBENZENE', 'n-NONYLBENZENE', 'ISOBUTYLBENZENE', 'n-UNDECYLBENZENE', '2-PHENYLBUTENE-1', 'o-ETHYLTOLUENE', 'm-DIVINYLBENZENE', 'm-ETHYLTOLUENE', 'p-ETHYLTOLUENE', 'cis-2-PHENYLBUTENE-2', 'trans-2-PHENYLBUTENE-2', 'n-PENTYLBENZENE', 'o-TERPHENYL', 'm-TERPHENYL', 'p-TERPHENYL', 'p-tert-BUTYLSTYRENE', '4-ISOBUTYLSTYRENE', 'p-tert-BUTYL ETHYLBENZENE', '1,2,4,5-TETRAMETHYLBENZENE', '1,2,3,5-TETRAMETHYLBENZENE', '1,2-DIPHENYLETHANE', 'trans-STILBENE', 'cis-STILBENE', 'm-DIISOPROPYLBENZENE', 'p-DIISOPROPYLBENZENE', 'TETRAPHENYLETHYLENE', 'n-TETRADECYLBENZENE', 'n-PENTADECYLBENZENE', 'n-HEXADECYLBENZENE', 'n-HEPTADECYLBENZENE', 'n-OCTADECYLBENZENE', 'BIPHENYL', 'ACENAPHTHENE', 'FLUORANTHENE', '1-n-BUTYLNAPHTHALENE', '1-n-PENTYLNAPHTHALENE', '1-PHENYLNAPHTHALENE', '1-ETHYLNAPHTHALENE', '2,6-DIETHYLNAPHTHALENE', '1-n-DECYLNAPHTHALENE', '2-ETHYLNAPHTHALENE', '1-n-HEXYLNAPHTHALENE', '1-METHYLNAPHTHALENE', '2-METHYLNAPHTHALENE', '2,6-DIMETHYLNAPHTHALENE', '2,7-DIMETHYLNAPHTHALENE', 'NAPHTHALENE', '2,3-DIMETHYL-2,3-DIPHENYLBUTANE', '1-n-NONYLNAPHTHALENE', '1-n-PROPYLNAPHTHALENE', '1-n-HEXYL-1,2,3,4-TETRAHYDRONAPHTHALENE', '2,4-DIPHENYL-4-METHYLPENTENE-1', 'TRIPHENYLMETHANE', 'TETRAPHENYLMETHANE', '1,1,2-TRIPHENYLETHANE', '1,1,2,2-TETRAPHENYLETHANE', 'ANTHRACENE', 'PHENANTHRENE', 'CHRYSENE', 'PYRENE', 'INDANE', 'ACENAPHTHALENE', 'FLUORENE', 'INDENE', '1-METHYLINDENE', '1,2,3-TRIMETHYLINDENE', '2-METHYLINDENE', '1,4-DI-tert-BUTYLBENZENE', '1,3,5-TRI-tert-BUTYLBENZENE', '1,3,5-TRIISOPROPYLBENZENE', 'PENTAMETHYLBENZENE', '1,2,4-TRIETHYLBENZENE', 'HEXAMETHYLBENZENE', '1,2,3-TRIETHYLBENZENE', '1,2,3,5-TETRAETHYLBENZENE', 'PENTAETHYLBENZENE', 'HEXAETHYLBENZENE', '1-(4-ETHYLPHENYL)-2-PHENYLETHANE', '1,2-DIMETHYL-3-PROPYLBENZENE', '1,2,3-TRIMETHYL-4-ETHYLBENZENE', '1,2,4-TRIMETHYL-3-ETHYLBENZENE', '1,2,4-TRIMETHYL-5-ETHYLBENZENE', '1-(4-ETHYLPHENYL)-2-(4-ETHYLPHENYL)ETHANE', 'o-ETHYLSTYRENE', 'm-ETHYLSTYRENE', 'p-ETHYLSTYRENE', '4-METHYLPHENANTHRENE', 'BENZANTHRACENE', 'NAPHTHACENE', '1-ETHYL-2-ISOPROPYLBENZENE', '1,3,5-TRIETHYLBENZENE', 'n-BUTYL MERCAPTAN', 'ETHYL MERCAPTAN', 'n-OCTYL MERCAPTAN', 'n-HEPTYL MERCAPTAN', 'n-HEXYL MERCAPTAN', 'METHYL MERCAPTAN', 'n-PENTYL MERCAPTAN', 'tert-BUTYL MERCAPTAN', 'sec-BUTYL MERCAPTAN', 'ISOPROPYL MERCAPTAN', 'n-PROPYL MERCAPTAN', 'tert-OCTYL MERCAPTAN', 'n-NONYL MERCAPTAN', 'n-DODECYL MERCAPTAN', 'ISOBUTYL MERCAPTAN', 'n-DECYL MERCAPTAN', 'UNDECYL MERCAPTAN', 'tert-DODECYL MERCAPTAN', 'PHENYL MERCAPTAN', 'CYCLOHEXYL MERCAPTAN', 'BENZYL MERCAPTAN', '1,2-ETHANEDITHIOL', 'CARBON DIOXIDE', 'HYDROGEN', 'HYDROGEN SULFIDE', 'NITROGEN', 'WATER', 'CO', 'Ar', 'O2']

# Temperature for all calculations
T = 298.15  # K

# Dictionary to store results
results = {"data": {}}

# Track successful fragmentations
successful_chemicals = {}

# First pass - identify all chemicals that fragment successfully
for cas in dippr_cmpds:
    try:
        chem = Chemical(cas)
        if chem.Tc is None or chem.Pc is None or chem.omega is None or chem.rdkitmol is None:
            continue
        assignment, _, _, success, _ = smarts_fragment_priority(
            catalog=EPPR78_GROUPS_LIST, 
            rdkitmol=chem.rdkitmol
        )
        
        if success:
            groups = {EPPR78_GROUPS_LIST[i-1].group: v for i, v in assignment.items()}
            successful_chemicals[chem.CAS] = {
                'groups': groups,
                'Tc': chem.Tc,
                'Pc': chem.Pc,
                'omega': chem.omega
            }
            
    except Exception as e:
        print(e)
        continue

# Second pass - calculate kij for all valid pairs
for (cas1, data1), (cas2, data2) in combinations(successful_chemicals.items(), 2):
    kij = PPR78_kij(
        T=T,
        molecule1_groups=data1['groups'],
        molecule2_groups=data2['groups'],
        Tc1=data1['Tc'],
        Pc1=data1['Pc'],
        omega1=data1['omega'],
        Tc2=data2['Tc'],
        Pc2=data2['Pc'],
        omega2=data2['omega'],
        version='extended'
    )

    # Store result in the format shown
    key = ' '.join(sorted_CAS_key([cas1, cas2]))
    results["data"][key] = {
        "kij": round(kij, 4),
    }
# Save results
output_path = os.path.join(os.path.dirname(__file__), '..', 'thermo', 'Interaction Parameters/eppr78_common.json')
os.makedirs(os.path.dirname(output_path), exist_ok=True)

with open(output_path, 'w') as f:
    json.dump(results, f)

print(f"Successfully processed {len(successful_chemicals)} chemicals")
print(f"Generated {len(results['data'])} interaction parameters")